// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once
#include <map>

#include "fastflow/framework/allocator.h"
#include "fastflow/framework/tensor.h"
#include "fastflow/graph/graph_transformer_mgr.h"
#include "fastflow/graph/insert_cast_transformer.h"

namespace fastflow {
class ExecutionProviders;
class Graph;
class GraphTransformerManager;
class InsertCastTransformer;
class KernelRegistryManager;
class NodeArg;
class SessionState;

namespace logging {
class Logger;
}

class SessionStateInitializer {
 public:
  SessionStateInitializer(fastflow::Graph& graph,
                          SessionState& session_state,
                          const ExecutionProviders& providers,
                          KernelRegistryManager& kernel_registry_manager,
                          const logging::Logger& logger);

  // First perform any transformations and create the execution plan
  Status CreatePlan(const fastflow::GraphTransformerManager& graph_transformation_manager,
                            const InsertCastTransformer& insert_cast_transformer,
                            const std::vector<const NodeArg*>& outer_scope_node_args,
                            bool enable_sequential_execution);

  // initialize tensors, and save. save kernels and input/output node mappings
  // @param enable_memory_pattern
  Status InitializeAndSave(bool enable_memory_pattern,
                                   std::map<FASTFLOWAllocatorInfo, BufferUniquePtr>& weights_buffers);

 private:
  fastflow::Graph& graph_;
  SessionState& session_state_;

  const ExecutionProviders& execution_providers_;
  KernelRegistryManager& kernel_registry_manager_;
  const logging::Logger& logger_;
};
}  // namespace fastflow
